import os
import cv2
import numpy as np
import torch
from mmengine.runner import load_checkpoint
from mmseg.apis import init_model, inference_model
from mmseg.structures import SegDataSample


# 颜色定义：BGR
PALETTE = {
    "lcx": (0, 255, 0),   # 绿色
    "lad": (0, 0, 255)    # 红色
}


def overlay_mask(image, mask, color, alpha=0.5):
    """在图像上叠加半透明的 mask."""
    overlay = image.copy()
    overlay[mask > 0] = color
    return cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)


def visualize_dual_branch(model, img_path, save_path):
    """对单张图像做推理并可视化结果."""
    result: SegDataSample = inference_model(model, img_path)

    image = cv2.imread(img_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 取出 LCX/LAD mask
    lcx_mask = result.pred_seg_map_head0.cpu().numpy()
    lad_mask = result.pred_seg_map_head1.cpu().numpy()

    # 叠加显示
    vis_img = overlay_mask(image, lcx_mask, PALETTE["lcx"], alpha=0.5)
    vis_img = overlay_mask(vis_img, lad_mask, PALETTE["lad"], alpha=0.5)

    # 保存结果
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    cv2.imwrite(save_path, cv2.cvtColor(vis_img, cv2.COLOR_RGB2BGR))
    print(f"可视化结果已保存到: {save_path}")


if __name__ == "__main__":
    config_file = "configs/main_vessel/resnet50_dualhead.py"  # ✅ 你的配置文件
    checkpoint_file = "work_dirs/resnet50_dualhead/latest.pth"  # ✅ 训练好的权重
    img_path = "dataset/images/validation/example.png"          # ✅ 待测试图像
    save_path = "vis_results/example_vis.png"

    model = init_model(config_file, checkpoint_file, device="cuda:0")
    visualize_dual_branch(model, img_path, save_path)
